import os
from openai import OpenAI

# 设置 DeepSeek API 密钥
DEEPSEEK_API_KEY = "sk-775bee8e690840c582899e92cf033bfa"
DEEPSEEK_API_URL = "https://api.deepseek.com"
MODEL = 'deepseek-reasoner'

# 自定义 LLM 类
class DeepSeekLLM():
    def __init__(self, api_key=DEEPSEEK_API_KEY, model = MODEL):
        self._api_key = api_key
        self._model = model


    def get(self, query, is_stream = False, stream_process_fun=None):
        client = OpenAI(api_key=self._api_key, base_url=DEEPSEEK_API_URL)
        response = client.chat.completions.create(
            model = self._model,
            messages=[
                {"role": "user", "content": query},
            ],
            stream=is_stream
        )
        str_list=[]
        if is_stream:
            for chunk in response:
                if chunk.choices[0].delta.content is not None:
                    content = chunk.choices[0].delta.content
                    if stream_process_fun:
                        stream_process_fun(content)
                    str_list.append(content)
        else:
            str_list.append(response.choices[0].message.content)

        return  "".join(str_list)

if __name__ == '__main__':
    # 使用自定义的 DeepSeek LLM
    llm = DeepSeekLLM()
    prompt = "介绍一下 Python 语言"
    response = llm.get(prompt, is_stream=True)
    print(response)